# Transport initialization module
#
# This module implements initialization logic from Fortran initialize.f90
# for setting up FLEXPART grid information and domain validation.

module TransportInit

using Dates
using ..FLEXINVERT
using ..Settings: Files, Config
using ..CoreTypes: Domain
using ..DatesUtil: juldate, caldate

# Include FlexpartIO module
include("flexpart.jl")
using .FlexpartIO

"""
    FlexpartGrids

Holds FLEXPART grid information for both global and nested domains.
This struct mirrors the global variables in Fortran's mod_var.f90.
"""
Base.@kwdef mutable struct FlexpartGrids
    # Global domain
    nxgrid::Int = 0
    nygrid::Int = 0
    llx::Float64 = 0.0
    lly::Float64 = 0.0
    dx::Float64 = 1.0
    dy::Float64 = 1.0
    nxshift::Int = 0
    nzgrid::Int = 0
    outheight::Vector{Float64} = Float64[]
    gbl_lon::Vector{Float64} = Float64[]
    gbl_lat::Vector{Float64} = Float64[]

    # Nested domain (if used)
    nnxgrid::Int = 0
    nnygrid::Int = 0
    nllx::Float64 = 0.0
    nlly::Float64 = 0.0
    ndx::Float64 = 1.0
    ndy::Float64 = 1.0
    nnxshift::Int = 0
    nest_lon::Vector{Float64} = Float64[]
    nest_lat::Vector{Float64} = Float64[]
    trajdays::Int = 0
end

"""
    DateInfo

Holds date and time calculation results from initialization.
"""
Base.@kwdef struct DateInfo
    juldatei::Float64 = 0.0  # Julian start date
    juldatef::Float64 = 0.0  # Julian end date
    nday::Int = 0            # Number of days in inversion interval
    ntstate::Int = 0         # Number of state time intervals
    ndt::Int = 1             # Number of time steps per 24h in prior
    nd_nee::Int = 0          # Number of NEE values per day (CO2 only)
    ntcini::Int = 0          # Number of initial mixing ratio time steps
end

"""
    find_flexpart_header(files::Files, config::Config, nested::Bool=false) -> String

Find a FLEXPART header file using the same logic as Fortran initialize.f90.
Searches through receptor directories for the appropriate month directory.

# Arguments
- `files`: Files configuration
- `config`: Configuration settings
- `nested`: Whether to look for nested header (header_nest) or global (header)

# Returns
- Path to the header file if found

# Throws
- `ErrorException` if no header file is found
"""
function find_flexpart_header(files::Files, config::Config, nested::Bool=false)
    yyyymm = config.datei ÷ 100
    adate = string(yyyymm, pad=6)

    # Read receptor list to get receptor names
    receptor_names = ["AAA"]  # Default from the test data

    if !isempty(files.file_recept) && isfile(files.file_recept)
        try
            # Extract receptor names from the receptor file
            # Format: name lat lon height (space-separated)
            lines = readlines(files.file_recept)
            unique_names = Set{String}()
            for line in lines
                line = strip(line)
                if !isempty(line) && !startswith(line, "#")
                    parts = split(line)
                    if length(parts) >= 1
                        push!(unique_names, parts[1])
                    end
                end
            end
            if !isempty(unique_names)
                receptor_names = collect(unique_names)
            end
        catch err
            @warn "Could not read receptor file, using default receptor name 'AAA'" err
        end
    end

    header_bases = nested ? ("header_nest", "header") : ("header",)

    # Resolve the flexpart path (handle relative paths)
    flexpart_path = if startswith(files.path_flexpart, "/")
        files.path_flexpart
    else
        abspath(files.path_flexpart)
    end

    for rec_name in receptor_names
        for base in header_bases
            # Prefer the human-readable text form when available
            filename_txt = joinpath(flexpart_path, rec_name, adate, "$(base)_txt")
            if isfile(filename_txt)
                if nested && base == "header"
                    @info "Nested header missing; falling back to global header text" filename = filename_txt
                end
                return filename_txt
            end

            # Otherwise test the binary FLEXPART header variant
            filename = joinpath(flexpart_path, rec_name, adate, base)
            if isfile(filename)
                if nested && base == "header"
                    @info "Nested header missing; falling back to global header" filename = filename
                end
                return filename
            end
        end
    end

    error("Cannot find FLEXPART $(nested ? "nested " : "")header for date $adate")
end

"""
    initialize_flexpart_grids!(grids::FlexpartGrids, files::Files, config::Config)

Initialize FLEXPART grid information from header files.
Implements the FLEXPART initialization logic from Fortran initialize.f90.

# Arguments
- `grids`: FlexpartGrids struct to populate
- `files`: Files configuration
- `config`: Configuration settings

# Notes
- Reads global header and optionally nested header if config.nested is true
- Calculates coordinate arrays (longitude/latitude)
- Handles global domain longitude shifts
"""
function initialize_flexpart_grids!(grids::FlexpartGrids, files::Files, config::Config)
    # Read global header
    global_header_path = find_flexpart_header(files, config, false)
    @info "Reading FLEXPART header: $global_header_path"

    # Determine if it's a text header or binary header
    if endswith(global_header_path, "_txt")
        global_header = FlexpartIO.read_header_text(global_header_path)
    else
        global_header = FlexpartIO.read_header(global_header_path)
    end

    # Copy global domain information
    grids.nxgrid = global_header.nxgrid
    grids.nygrid = global_header.nygrid
    grids.llx = global_header.llx
    grids.lly = global_header.lly
    grids.dx = global_header.dx
    grids.dy = global_header.dy
    grids.nxshift = global_header.nxshift
    grids.nzgrid = global_header.nzgrid
    grids.outheight = copy(global_header.outheight)
    grids.trajdays = global_header.trajdays

    # Calculate global coordinate arrays
    grids.gbl_lon = [grids.llx + (i-1)*grids.dx for i = 1:grids.nxgrid]
    grids.gbl_lat = [grids.lly + (i-1)*grids.dy for i = 1:grids.nygrid]

    # Read nested header if needed
    if config.nested
        nested_header_path = find_flexpart_header(files, config, true)
        @info "Reading FLEXPART nested header: $nested_header_path"

        if endswith(nested_header_path, "_txt")
            nested_header = FlexpartIO.read_header_text(nested_header_path)
        else
            nested_header = FlexpartIO.read_header(nested_header_path)
        end

        # Copy nested domain information
        grids.nnxgrid = nested_header.nxgrid
        grids.nnygrid = nested_header.nygrid
        if nested_header.trajdays > 0
            grids.trajdays = nested_header.trajdays
        end
        grids.nllx = nested_header.llx
        grids.nlly = nested_header.lly
        grids.ndx = nested_header.dx
        grids.ndy = nested_header.dy
        grids.nnxshift = nested_header.nxshift

        # Calculate nested coordinate arrays
        grids.nest_lon = [grids.nllx + (i-1)*grids.ndx for i = 1:grids.nnxgrid]
        grids.nest_lat = [grids.nlly + (i-1)*grids.ndy for i = 1:grids.nnygrid]
    end
end

"""
    initialize_dates(config::Config) -> DateInfo

Initialize date and time variables for the inversion.
Implements date calculation logic from Fortran initialize.f90.

# Arguments
- `config`: Configuration settings

# Returns
- `DateInfo`: Struct with calculated date and time information
"""
function initialize_dates(config::Config)
    # Convert to Date objects for calculation
    datei_str = string(config.datei)
    datef_str = string(config.datef)

    if length(datei_str) != 8 || length(datef_str) != 8
        error("Invalid date format. Expected yyyymmdd, got datei=$(config.datei), datef=$(config.datef)")
    end

    # Parse dates
    year_i = parse(Int, datei_str[1:4])
    month_i = parse(Int, datei_str[5:6])
    day_i = parse(Int, datei_str[7:8])

    year_f = parse(Int, datef_str[1:4])
    month_f = parse(Int, datef_str[5:6])
    day_f = parse(Int, datef_str[7:8])

    date_i = Date(year_i, month_i, day_i)
    date_f = Date(year_f, month_f, day_f)

    # Calculate Julian dates using DatesUtil (Fortran-compatible)
    juldatei = juldate(config.datei, 0)  # Start of day (00:00:00)
    juldatef = juldate(config.datef, 0)  # Start of day (00:00:00)

    # Number of days in inversion interval
    nday = Int(juldatef - juldatei + 1)

    # Initialize time step variables based on species
    if config.spec == "co2"
        # CO2-specific calculations
        if config.nstep_nee_reg == 0
            error("nstep_nee_reg cannot be zero for CO2")
        end
        nd_nee = 24 ÷ config.nstep_nee_reg

        # State variable time intervals
        statres = config.statres
        statres_hr = config.statres_hr
        ndt = Int(24.0 / statres_hr)
        ntstate = Int(Float64(nday) / statres)
    else
        # Non-CO2 species
        nd_nee = 0
        statres = config.statres
        ntstate = Int(Float64(nday) / statres)
        ndt = 1
    end

    if ntstate < 1
        error("State time steps < 1. Increase time interval (statres = $(config.statres))")
    end

    # Adjust nday to match number of state time steps
    nday_adjusted = ntstate * Int(config.statres)
    juldatef_adjusted = juldatei + Float64(nday_adjusted) - 1.0

    if nday_adjusted != nday
        @info "Adjusted nday from $nday to $nday_adjusted to match state time steps"
        # Convert back to calendar date for logging
        yyyymmdd_adjusted, hhmmss_adjusted = caldate(juldatef_adjusted)
        @info "Adjusted end date to $(yyyymmdd_adjusted)"
    end

    # Number of initial mixing ratio time steps
    ntcini = max(1, nday_adjusted ÷ config.cinires)

    return DateInfo(
        juldatei = juldatei,
        juldatef = juldatef_adjusted,
        nday = nday_adjusted,
        ntstate = ntstate,
        ndt = ndt,
        nd_nee = nd_nee,
        ntcini = ntcini
    )
end

"""
    validate_domain_bounds(domain::Domain, grids::FlexpartGrids, config::Config)

Validate that the inversion domain is within FLEXPART domain bounds.
Implements boundary checking logic from Fortran initialize.f90.

# Arguments
- `domain`: Regional inversion domain
- `grids`: FLEXPART grid information
- `config`: Configuration settings

# Throws
- `ErrorException` if domain boundaries are outside FLEXPART domains
"""
function validate_domain_bounds(domain::Domain, grids::FlexpartGrids, config::Config)
    # Check against global domain
    if domain.rllx < grids.llx
        error("Western boundary ($(domain.rllx)) outside FLEXPART global domain ($(grids.llx))")
    end

    if domain.rurx > (grids.llx + grids.nxgrid * grids.dx)
        error("Eastern boundary ($(domain.rurx)) outside FLEXPART global domain ($(grids.llx + grids.nxgrid * grids.dx))")
    end

    if domain.rlly < grids.lly
        error("Southern boundary ($(domain.rlly)) outside FLEXPART global domain ($(grids.lly))")
    end

    if domain.rury > (grids.lly + grids.nygrid * grids.dy)
        error("Northern boundary ($(domain.rury)) outside FLEXPART global domain ($(grids.lly + grids.nygrid * grids.dy))")
    end

    # Check resolution match for global domain
    if abs(grids.dx - domain.rdx) > 1e-6 || abs(grids.dy - domain.rdy) > 1e-6
        error("FLEXPART global resolution ($(grids.dx)°×$(grids.dy)°) does not match inversion resolution ($(domain.rdx)°×$(domain.rdy)°)")
    end

    # Additional checks for nested domain
    if config.nested
        if domain.rllx < grids.nllx
            error("Western boundary ($(domain.rllx)) outside FLEXPART nested domain ($(grids.nllx))")
        end

        if domain.rurx > (grids.nllx + grids.nnxgrid * grids.ndx)
            error("Eastern boundary ($(domain.rurx)) outside FLEXPART nested domain ($(grids.nllx + grids.nnxgrid * grids.ndx))")
        end

        if domain.rlly < grids.nlly
            error("Southern boundary ($(domain.rlly)) outside FLEXPART nested domain ($(grids.nlly))")
        end

        if domain.rury > (grids.nlly + grids.nnygrid * grids.ndy)
            error("Northern boundary ($(domain.rury)) outside FLEXPART nested domain ($(grids.nlly + grids.nnygrid * grids.ndy))")
        end

        # Check resolution match for nested domain
        if abs(grids.ndx - domain.rdx) > 1e-6 || abs(grids.ndy - domain.rdy) > 1e-6
            error("FLEXPART nested resolution ($(grids.ndx)°×$(grids.ndy)°) does not match inversion resolution ($(domain.rdx)°×$(domain.rdy)°)")
        end
    end

    @info "Domain bounds validation passed"
end

"""
    initialize_transport_system(files::Files, config::Config, domain::Domain) -> (FlexpartGrids, DateInfo)

Main initialization function that sets up the transport system.
Implements the core logic from Fortran initialize.f90.

# Arguments
- `files`: Files configuration
- `config`: Configuration settings
- `domain`: Regional inversion domain

# Returns
- `FlexpartGrids`: FLEXPART grid information
- `DateInfo`: Date and time calculation results

# Notes
This function:
1. Initializes FLEXPART grids from header files
2. Calculates date and time variables
3. Validates domain boundaries
4. Performs consistency checks
"""
function initialize_transport_system(files::Files, config::Config, domain::Domain)
    @info "Initializing transport system..."

    # Initialize FLEXPART grids
    grids = FlexpartGrids()
    initialize_flexpart_grids!(grids, files, config)

    if !isempty(grids.outheight)
        domain.surface_layer_height = grids.outheight[1]
        domain.outheight = copy(grids.outheight)
    end

    domain.global_nxgrid = grids.nxgrid
    domain.global_nygrid = grids.nygrid
    domain.global_nzgrid = grids.nzgrid
    domain.global_nxshift = grids.nxshift
    domain.global_lon = copy(grids.gbl_lon)
    domain.global_lat = copy(grids.gbl_lat)
    if grids.trajdays > 0
        domain.trajdays = grids.trajdays
    end

    # Initialize date calculations
    date_info = initialize_dates(config)

    @info "Date calculations:"
    @info "  nday: $(date_info.nday)"
    @info "  statres: $(config.statres)"
    @info "  ntstate: $(date_info.ntstate)"
    if config.spec == "co2"
        @info "  nd_nee: $(date_info.nd_nee)"
        @info "  ndt: $(date_info.ndt)"
    end

    @info "Regional domain:"
    @info "  w_edge_lon: $(domain.rllx)"
    @info "  s_edge_lat: $(domain.rlly)"
    @info "  e_edge_lon: $(domain.rurx)"
    @info "  n_edge_lat: $(domain.rury)"
    @info "  xres: $(domain.rdx)"
    @info "  yres: $(domain.rdy)"
    @info "  nxregrid: $(domain.nxregrid)"
    @info "  nyregrid: $(domain.nyregrid)"

    @info "FLEXPART global domain:"
    @info "  llx: $(grids.llx)"
    @info "  lly: $(grids.lly)"
    @info "  dx: $(grids.dx)"
    @info "  dy: $(grids.dy)"
    @info "  nxgrid: $(grids.nxgrid)"
    @info "  nygrid: $(grids.nygrid)"
    @info "  nzgrid: $(grids.nzgrid)"

    if config.nested
        @info "FLEXPART nested domain:"
        @info "  nllx: $(grids.nllx)"
        @info "  nlly: $(grids.nlly)"
        @info "  ndx: $(grids.ndx)"
        @info "  ndy: $(grids.ndy)"
        @info "  nnxgrid: $(grids.nnxgrid)"
        @info "  nnygrid: $(grids.nnygrid)"
    end

    # Validate domain bounds
    validate_domain_bounds(domain, grids, config)

    @info "Transport system initialization completed successfully"

    return grids, date_info
end

end # module
